import torch
import torch.onnx
import torch.jit

from my_classifier import LeNet

input = torch.rand(1, 1, 32, 32, dtype=torch.float32)
model = LeNet()
input_names = ["input"]
output_names = ["output"]

torch.onnx.export(
    model, input, "leNet.onnx", verbose=True, export_params=True,
    input_names=input_names, output_names=output_names
)

# m = torch.jit.load("PicProc/partTypeClassifier.pt")

# print(m)

